from colorama import Fore
import json
import os
import time
import random

random.seed(0)


def handle(train_data, gold_data, size=10):

    output_dir = f'files/results/{train_data.split("/")[-1].split(".")[0]}'
    real_dir = f"media/{output_dir}"
    if not os.path.exists(real_dir):
        os.makedirs(real_dir)
    data_path = "media/data"
    idx = random.randint(1, 5)
    if size == 10:
        loss_info_file = (
            f"{data_path}/Size {size}/Net-{idx}/faster_40_50_d3_mode_save_info.json"
        )
        result_file = f"{data_path}/result_info_{size}.json"
    elif size == 100:
        loss_info_file = (
            f"{data_path}/Size{size}/Net-{idx}/faster_20_200_d3_mode_save_info.json"
        )
        result_file = f"{data_path}/result_info_{size}.json"

    # 将loss_info_file直接复制到result_file
    with open(loss_info_file, "r") as f:
        data = json.load(f)
    training_loss = f"{real_dir}/training_loss.json"
    with open(training_loss, "w") as f:
        json.dump(data, f, indent=4)

    with open(result_file, "r") as f:
        data = json.load(f)
    # 需要做数据抽取
    result_info_file = f"{real_dir}/result_info.json"
    result_info = {}
    result_info["gold_data"] = data["gold_data"][f"Net-{idx}"]
    result_info["train_data"] = data["train_data"][f"Net-{idx}"]
    result_info["train_metrics"] = data["train_metrics"][f"Net-{idx}"]
    with open(result_info_file, "w") as f:
        json.dump(result_info, f, indent=4)

    return output_dir


if __name__ == "__main__":
    train_data = "files/2fc6eb46e6.tsv"
    gold_data = "files/c43f36953a.tsv"
    output_dir = "files/results/2fc6eb46e6"
    handle(train_data, gold_data)
# ---- END
